#version 450
#extension GL_ARB_separate_shader_objects : enable
#extension GL_EXT_nonuniform_qualifier : enable
#extension GL_EXT_scalar_block_layout : enable
#extension GL_GOOGLE_include_directive : enable
#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require
#extension GL_EXT_buffer_reference2 : require
#extension GL_ARB_shader_image_size : require

#include "Common.glsl"

layout (location = 0) in vec2 inUV;

layout (binding = 0) uniform Ubo {
    mat4 projview;
	mat4 viewInverse;
	mat4 projInverse;
    uint frame;
} ubo;

layout (binding = 1, rgba32f) uniform image2D resultImage;
layout (binding = 2) uniform sampler2D samplerHdr;
layout (binding = 3) uniform sampler2D samplerCache;

vec3 SampleHdr(float r1, float r2) {
    vec2 xy = texture(samplerCache, vec2(r1, r2)).rg; // x, y
    xy.y = 1.0 - xy.y; // flip y

    // 获取角度
    float phi = 2.0 * M_PI * (xy.x - 0.5);    // [-pi ~ pi]
    float theta = M_PI * (xy.y - 0.5);        // [-pi/2 ~ pi/2]   

    // 球坐标计算方向
    vec3 L = vec3(cos(theta)*cos(phi), sin(theta), cos(theta)*sin(phi));

    return L;
}

// 将三维向量 v 转为 HDR map 的纹理坐标 uv

vec2 ToSphericalCoord(vec3 v) {
    vec2 uv = vec2(atan(v.z, v.x), asin(v.y));
    uv *= vec2(0.1591, 0.3183);
    uv += 0.5;
    uv.y = 1.0 - uv.y;
    return uv;
}

// 获取 HDR 环境颜色
vec3 HdrColor(vec3 L) {
    vec2 uv = ToSphericalCoord(normalize(L));
    vec3 color = texture(samplerHdr, uv).rgb;
    return color;
}

// 输入光线方向 L 获取 HDR 在该位置的概率密度
// hdr 分辨率为 4096 x 2048 --> hdrResolution = 4096
float HdrPdf(vec3 L) {
    vec2 uv = ToSphericalCoord(normalize(L));   // 方向向量转 uv 纹理坐标

    float pdf = texture(samplerCache, uv).b;      // 采样概率密度
    float theta = M_PI * (0.5 - uv.y);            // theta 范围 [-pi/2 ~ pi/2]
    float sin_theta = max(sin(theta), 1e-10);

    // 球坐标和图片积分域的转换系数
    float p_convert = float(2048 * 1024) / (2.0 * M_PI * M_PI * sin_theta);  
    
    return pdf * p_convert;
}

void main()
{
    ivec2 texUV = imageSize(resultImage);
    texUV.x = int(inUV.x * texUV.x);
    texUV.y = int(inUV.y * texUV.y);
    uint seed = Tea(texUV.x * texUV.y + texUV.y, ubo.frame);

    vec2 uv = inUV * 2.0 - 1.0;
	vec3 origin = vec3(ubo.viewInverse * vec4(0, 0, 0, 1));
    vec3 target = vec3(ubo.projInverse * vec4(uv.x, uv.y, 1, 1));
    vec3 direction = vec3(ubo.viewInverse * vec4(normalize(target), 0));


    vec3 color = vec3(0.0);

    float rand1 = Random(seed);
    float rand2 = Random(seed);

#if 0

    // vec3 L = SampleHdr(inUV.x, inUV.y);
    // vec3 L = SampleHdr(rand1, rand2);
    // color = HdrColor(L);
    // float pdfLight = HdrPdf(L);
    // color = vec3(pdfLight / (2048.0 * 1024.0));

    // vec2 hdrUV = ToSphericalCoord(direction);

    // 这么写才等价于上面那条代码
    // direction.y *= -1.0;
    // vec2 hdrUV = SampleSphericalMap(direction);

    // color = texture(samplerHdr, hdrUV).rgb;

    // float pdfLight = HdrPdf(direction);
    // color = HdrColor(direction);
    // color = vec3(pdfLight / (2048.0 * 1024.0));

    // vec2 hdrUV = ToSphericalCoord(direction);


    // 球坐标和图片积分域的转换系数, 把概率转换回 0 - 1
    float p_convert = 1.0 / (2.0 * M_PI * M_PI);
    float pdf = texture(samplerCache, inUV).b;
    color = vec3(pdf * 2048.0 * 1024.0 * p_convert);
    

    
#else

    color = vec3(texture(samplerCache, inUV).b) * 10000.0;

    vec2 cacheUV = texture(samplerCache, vec2(rand1, rand2)).rg;
    if(abs(inUV.x - cacheUV.x) < 0.03 && abs(inUV.y - cacheUV.y) < 0.03)
    {
        color = vec3(1.0, 0.0, 0.0);
    }

#endif


#if 1
    if(ubo.frame > 0)
    {
        float a = 1.0f / float(ubo.frame + 1);
        vec3  oldColor = imageLoad(resultImage, texUV).xyz;
        imageStore(resultImage, texUV, vec4(mix(oldColor, color, a), 1.f));
    }
    else
    {
        imageStore(resultImage, texUV, vec4(color, 1.f));
    }
#else
    imageStore(resultImage, texUV, vec4(color, 1.f));
#endif
}